#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from cv_bridge import CvBridge, CvBridgeError
import rospy
import cv2
import numpy as np
import rospkg
from sensor_msgs.msg import Image
from std_msgs.msg import String
import threading
import os

class ColorDetector:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('multi_template_matcher', anonymous=False)
        
        # 初始化CV Bridge
        self.bridge = CvBridge()
        
        # 算法参数配置
        self.target_fps = rospy.get_param('~target_fps', 30)
        self.downscale_resolution = rospy.get_param('~downscale_resolution', True)
        self.box_size = rospy.get_param('~box_size', 200)  # 中心框大小
        self.confidence_threshold = 0.1  # 置信度阈值
        
        # 定义颜色范围（HSV格式）[6,7](@ref)
        self.color_ranges = {
            'red': ([0, 100, 100], [10, 255, 255], [160, 100, 100], [180, 255, 255]),
            'green': ([35, 50, 50], [85, 255, 255]),
            'black': ([0, 0, 0], [180, 255, 50])
        }
        
        # 图像缓冲区和处理状态标志
        self.latest_image = None
        self.image_lock = threading.Lock()
        self.processing_flag = False
        
        # 创建识别结果发布者
        self.color_pub = rospy.Publisher('~detected_color', String, queue_size=10)
        
        # 初始化订阅者
        self.image_sub = rospy.Subscriber(
            '/camera/rgb/image_raw', 
            Image, 
            self.image_callback,
            queue_size=1,
            buff_size=2**24,
            tcp_nodelay=True
        )
        
        # 初始化显示窗口
        cv2.namedWindow('Color Detection', cv2.WINDOW_NORMAL)
        rospy.loginfo(f"颜色检测节点已启动 | 中心框大小: {self.box_size}px | 置信度阈值: {self.confidence_threshold}")
        
        # 创建处理线程
        self.processing_thread = threading.Thread(target=self.processing_loop)
        self.processing_thread.daemon = True
        self.processing_thread.start()
        
        # 注册关闭回调
        rospy.on_shutdown(self.shutdown)

    def detect_color(self, hsv_img):
        """检测给定HSV图像中的主要颜色"""
        # 创建颜色掩码并计算像素比例
        color_pixels = {}
        
        for color_name, ranges in self.color_ranges.items():
            if color_name == 'red':  # 红色有两个范围[6](@ref)
                lower1 = np.array(ranges[0], dtype="uint8")
                upper1 = np.array(ranges[1], dtype="uint8")
                lower2 = np.array(ranges[2], dtype="uint8")
                upper2 = np.array(ranges[3], dtype="uint8")
                
                mask1 = cv2.inRange(hsv_img, lower1, upper1)
                mask2 = cv2.inRange(hsv_img, lower2, upper2)
                mask = cv2.bitwise_or(mask1, mask2)
            else:
                lower = np.array(ranges[0], dtype="uint8")
                upper = np.array(ranges[1], dtype="uint8")
                mask = cv2.inRange(hsv_img, lower, upper)
            
            # 计算该颜色在ROI中的比例
            pixel_count = cv2.countNonZero(mask)
            total_pixels = hsv_img.shape[0] * hsv_img.shape[1]
            color_pixels[color_name] = pixel_count / total_pixels if total_pixels > 0 else 0
        
        # 找出占比最大的颜色
        detected_color = max(color_pixels, key=color_pixels.get)
        confidence = color_pixels[detected_color]
        
        # 如果黑色占比高但亮度也高，则可能不是真正的黑色
        if detected_color == 'black' and np.mean(hsv_img[:,:,2]) > 100:
            # 寻找次高占比的颜色
            color_pixels.pop('black')
            if color_pixels:
                detected_color = max(color_pixels, key=color_pixels.get)
                confidence = color_pixels[detected_color]
        
        return detected_color, confidence

    def image_callback(self, msg):
        """快速图像回调，仅存储最新图像"""
        with self.image_lock:
            self.latest_image = msg

    def processing_loop(self):
        """独立处理线程，从缓冲区获取图像进行处理"""
        rate = rospy.Rate(self.target_fps)
        while not rospy.is_shutdown():
            with self.image_lock:
                if self.latest_image is not None and not self.processing_flag:
                    self.processing_flag = True
                    image_msg = self.latest_image
                    self.latest_image = None
                    
                    try:
                        # 转换ROS图像消息
                        cv_image = self.bridge.imgmsg_to_cv2(image_msg, "bgr8")
                        
                        # 图像降采样
                        if self.downscale_resolution:
                            height, width = cv_image.shape[:2]
                            if width > 640 or height > 480:
                                cv_image = cv2.resize(cv_image, (640, 480))
                        
                        # 获取图像中心区域（缩小为原来的80%）
                        height, width = cv_image.shape[:2]
                        center_x, center_y = width // 2, height // 2
                        box_half = int(self.box_size * 0.4)  # 缩小中心区域
                        
                        # 计算检测区域坐标
                        x1 = max(0, center_x - box_half)
                        y1 = max(0, center_y - box_half)
                        x2 = min(width, center_x + box_half)
                        y2 = min(height, center_y + box_half)
                        
                        # 提取中心区域ROI
                        roi = cv_image[y1:y2, x1:x2]
                        
                        if roi.size == 0:
                            rospy.logwarn("中心区域太小，无法检测")
                            continue
                        
                        # 转换为HSV颜色空间[6,7](@ref)
                        hsv_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
                        
                        # 检测主要颜色
                        detected_color, confidence = self.detect_color(hsv_roi)
                        
                        # 仅在置信度高于阈值时发布结果
                        if confidence >= self.confidence_threshold:
                            # 发布检测到的颜色
                            self.color_pub.publish(detected_color)
                            rospy.loginfo(f"检测到颜色: {detected_color} (置信度: {confidence:.2f})")
                            
                            # 在图像上绘制中心框和结果
                            cv2.rectangle(cv_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            
                            # 根据检测到的颜色设置文本颜色
                            text_color = (0, 0, 255) if detected_color == 'red' else \
                                        (0, 255, 0) if detected_color == 'green' else \
                                        (0, 0, 0)
                            
                            cv2.putText(cv_image, f"{detected_color} ({confidence:.2f})", 
                                       (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
                                       0.7, text_color, 2)
                        else:
                            # 置信度不足时，绘制框但不显示颜色
                            cv2.rectangle(cv_image, (x1, y1), (x2, y2), (255, 255, 0), 2)
                            cv2.putText(cv_image, "Low confidence", 
                                       (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 
                                       0.7, (255, 255, 0), 2)
                        
                        # 显示结果
                        cv2.imshow('Color Detection', cv_image)
                        cv2.waitKey(1)
                    except CvBridgeError as e:
                        rospy.logerr(f"图像转换失败: {str(e)}")
                    except Exception as e:
                        rospy.logerr(f"处理异常: {str(e)}")
                    finally:
                        self.processing_flag = False
            
            rate.sleep()

    def shutdown(self):
        """节点关闭时的清理操作"""
        cv2.destroyAllWindows()
        rospy.loginfo("节点已安全关闭")

if __name__ == '__main__':
    detector = ColorDetector()
    rospy.spin()
